"""
This is the Sierra Nevada example from:

B. Avdeev, N. A. Niemi, M. K. Clark, 2011. 
Doing more with less: Bayesian estimation of erosion models with detrital thermochronometric data. 
EPSL, doi:10.1016/j.epsl.2011.03.020
"""



__author__="Boris Avdeev, borisaqua@gmail.com"

import sys
import pymc as pm
import numpy as np
import matplotlib.pyplot as plt
import common as ba

def InitExhumation(brakes):
    """Initialize piece-wise linear exhumation model"""
    Parm = []
    for i in range(1,brakes+2):
        Parm.append(pm.Uniform("e%i" % i, 0.001, 1, value=0.05))
    Parm.append(pm.Uniform('hc', 0., 3.))
    abr = 0
    for i in range(1,brakes+1):
        abr = pm.Uniform("abr%i" % i, abr, 60)
        Parm.append(abr)
    return Parm

class DetritalModel:
    """Set up a simple detrital thermochron model"""
    def __init__(self, n_abr):
        hyps = bins['h']
        hyps_w = bins['w']
        ages = al_age
        self.err = pm.Uniform('RelErr',0.1,0.2)
        self.parm = InitExhumation(n_abr)
        # Value p is probability of sampling from different hypsometry bins.
        # Here it is set to the weights of elevation bins, i.e. assuming uniform sampling.
        # To model nonuniform sampling, set p to Dirichlet distribution.
        self.idx = pm.Categorical('Index', p = hyps_w, size=len(ages))
        self.exp = pm.Lambda("TrueAHe", lambda parm=self.parm, hyps=hyps, idx=self.idx: ba.h2a(hyps[idx],parm))
        self.obs = pm.Normal("ObsAHe", mu = self.exp, tau = 1./(self.err*self.exp)**2, value = ages, observed=True)
        self.sim = pm.Lambda("SimAHe", lambda ta=self.exp, err=self.err: pm.rnormal(mu = ta, tau = 1./(err*ta)**2))

def plot(M,chain=-1):
    fig = plt.figure()
    fig.set_size_inches(6.,10.)
    p1 = plt.subplot(211,title='AHe')
    ba.gof(al_age,M.trace("SimAHe",chain=chain)[:])
    
    p2 = plt.subplot(212,sharex=p1)
    parms = np.vstack(par.trace(chain=chain) for par in M.parm).transpose()
    ba.ah(parms,4)
    
    p2.set_xlim(0,65)
    p2.set_ylim(0,4)
    fig.savefig("summary.png") 
    
    
                    

if __name__ == '__main__':
    """First argument is number of breaks, second is length of the chain."""
    bins = ba.read_xyh("data/Inyo.xyhs")
    al_age = np.genfromtxt("data/InyoAges.txt")
    
    n_brk = int(sys.argv[1])
    N = int(sys.argv[2])
    burn = N/5
    thin = (N-burn) / 1000
    name = "Inyo_%ibrk" % n_brk
    attempt = 0
    model=None
    while attempt<5000:
        try:
            model = DetritalModel(n_brk)
            break
        except pm.ZeroProbability:
            attempt+=1
            print "Init failure %i" % attempt

    db = pm.database.sqlite.load(name + '.sqlite')
    M = pm.MCMC(model, db=db, name=name)
    ##This might make a better chain (or might not).
    #M.use_step_method(pm.AdaptiveMetropolis, M.parm)
    M.sample(N,burn=burn,thin=thin)
    #pm.Matplot.plot(M)
    #pm.Matplot.summary_plot(M)
    plot(M)    
    M.db.close()


